#!/usr/bin/python
# -*- coding: utf-8 -*-
from pandas import DataFrame
from numpy import nan

from core.index_base import IndexBase


class BOLL(IndexBase):
    """
    Bollinger Band 布林带通道
    """
    def __init__(self, period=20, times=2, key='close'):
        super(BOLL, self).__init__()
        self.period = period
        self.key = key
        self.times = times

    def compute(self, data:DataFrame):
        boll_upbound_key = 'boll_upbound' + str(self.period)
        boll_downbound_key = 'boll_downbound' + str(self.period)
        boll_midbound_key = 'boll_midbound' + str(self.period)
        self.added_keys.extend([boll_upbound_key, boll_downbound_key, boll_midbound_key])

        bool_std_key = 'bool_std'
        data[boll_upbound_key] = nan
        data[boll_downbound_key] = nan
        data[boll_midbound_key] = nan
        data[bool_std_key] = nan
        data[boll_midbound_key] = data[self.key].rolling(self.period).mean()
        data[bool_std_key] = data[self.key].rolling(self.period).std()
        data[boll_upbound_key] = data[boll_midbound_key] + self.times * data[bool_std_key]
        data[boll_downbound_key] = data[boll_midbound_key] - self.times * data[bool_std_key]
        data = data.drop(labels=bool_std_key, axis=1)


if __name__=="__main__":
    from domain.transaction_data.repository import transaction_data_repository

    BOLL(period=20).compute(transaction_data_repository.get_contain_latest_data_by_tushare('002044.SZ'))